clear;clc;
tic;
%% set parameters
Nx=14;
Ny=14; % must be even
Nxa=Nx-1;
Nya=Ny-1;
N1=Nx*Ny; % on-lattice
N2=(Nx-1)*(Ny-1); % off-lattice
N=N1+N2;

VBB=1;
VAB=2.628*VBB;
V=VBB/2;
v=VAB-2*VBB;

Omega1=0.65*VBB; % final Rabi freq., subscript 1 for on-lattice atom B with n=113
Omega2=Omega1*(76/113)^2; % Rabi freq., subscript 2 for off-lattice atom A with n=76
Delta1=1*VBB; % atom B
% delta=0.1*Omega1;
Delta2=2*Delta1; % atom A
mu1=-Delta1;
mu2=-Delta2;
V1=VBB; % BB
V2=VAB; % AB
J1=Omega1^2*V1/(4*Delta1*(Delta1+V1));
J2=Omega2^2*V2/(4*Delta2*(Delta2+V2));

% Hamiltonian setting
% % Js
HJs0=Js(Nx,Ny,0,0);
% % mus
Hmus0=mus(Nx,Ny,mu1,mu2);

%% set Hamiltonian and initial state
H0=HJs0+Hmus0;

psi=zeros(N,1);

% initial configuration

% 
psi=ones(N,1);

% normalization
psi=psi/sqrt(sum(abs(psi).^2));

% collect coordinates
positions=Lattice(Nx,Ny);

%% pre-plotting
figure
hold on
axis equal
axis off
colormap("sky");
cb=colorbar;
cb.Label.String='Norm. Rydberg state population';
cb.Label.FontWeight='bold';
cb.Label.FontSize=20;
xpoints=zeros(1,N);
ypoints=zeros(1,N);
xpoints_main=zeros(1,N1);
ypoints_main=zeros(1,N1);
for jj=1:N
    xpoints(jj)=positions{jj}(1);
    ypoints(jj)=positions{jj}(2);
end

for jj=1:N1
    xpoints_main(jj)=positions{jj}(1);
    ypoints_main(jj)=positions{jj}(2);
end
ylim([mean(ypoints_main)+.96*(min(ypoints_main)-mean(ypoints_main)),...
    mean(ypoints_main)+.96*(max(ypoints_main)-mean(ypoints_main))])
tin=delaunayTriangulation([xpoints_main.' ypoints_main.']);
triplot(tin,'-','LineWidth',2,'Color',[0.15 0.3 0.6]);
MarkScale=140;
hfg=scatter(xpoints,ypoints,MarkScale,'MarkerEdgeColor',[0.15 0.3 0.6],'LineWidth',2.4);

for kk=1:N
    Prr0{kk}(1)=abs(psi(kk)).^2;
    scatter(xpoints(kk),ypoints(kk),0.9*MarkScale,Prr0{kk}(1),"filled")
end

set(gca,'FontSize',20)
set(gcf,'position',[300,0,1000,1000])

%% evolution
time_step=1e5;
t_final=1e5;
dt=t_final/time_step;
times=0:dt:t_final;

fidelity=1;

Prrs=cell(1,N);
for jj=1:N
    Prrs{jj}=zeros(1,length(times));
end

Eg=zeros(1,time_step);
Ee=zeros(1,time_step);
Ene=zeros(1,time_step);

for jj=1:length(times)
    t=times(jj);
    latexf=['Ω / V = 1.3',' (t = ',num2str(t),')'];
    lgh=title(latexf);
    set(lgh,'FontWeight','bold');
    Omega1t=t/t_final*Omega1;
    Omega2t=2*Omega1t;
    J1t=Omega1t^2*V1/(4*Delta1*(Delta1+V1));
    J2t=Omega2t^2*V2/(4*Delta2*(Delta2+V2));
    
    Jst=Js(Nx,Ny,J1t,J2t);
    must=mus(Nx,Ny,mu1,mu2);
    Ht=Jst+must;
    [D,V]=eig(Ht);
    enes=sort(diag(V));
    Eg(jj)=enes(1);
    Ee(jj)=enes(2);
    Ene(jj)=psi'*Ht*psi;
    for kk=1:N
        Prrs{kk}(jj)=abs(psi(kk)).^2;
    end
    psi_last=psi;
    psi=expm(-1i*Ht*dt)*psi;
    fidelity=[fidelity abs(psi'*psi_last).^2];
end

%
    for kk=1:N
        Prrs{kk}(length(times))=abs(psi(kk)).^2;       
        hgf=scatter(xpoints(kk),ypoints(kk),0.9*MarkScale,Prrs{kk}(length(times)),"filled");
    end
%
toc;

%% correlation fuction
Cij=zeros(N1);
for ii=1:N1
    for jj=1:N1
        Cij(ii,jj)=abs(psi(ii)*psi(jj));
    end
end
iis=1:N1;
jjs=1:N1;

figure
axis equal
box on
imagesc(iis,jjs,Cij);
colormap("sky");
cb=colorbar;
cb.Label.String='Correlation function C(i,j) of target atoms';
cb.Label.FontWeight='bold';
cb.Label.FontSize=20;
set(gca,'FontSize',20)
set(gcf,'position',[300,0,1000,1000])
latexf=['Ω / V = 1.3',' (t = 10000)'];
lgh=title(latexf);
set(lgh,'FontWeight','bold');
savefig('con2C22.fig')